package com.saascloud.thread;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import com.alibaba.ttl.threadpool.TtlExecutors;
import org.slf4j.MDC;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.*;
import java.util.function.Supplier;

/**
 * @Author: liangxu
 * @Contact: one_dream@aliyun.com
 * @Date: 2022/5/25 9:27
 * @Version: 1.0
 * @Description: 包装原始的runnable和callable，用于上下文传递
 * 所有返回的ExecutorService都支持租户上下文传递
 */
public class ThreadTaskUtils {

    public static final int DEFAULT_WORK_QUEUE_CAPACITY = 1024;

    public static void main(String[] args) {
        ThreadLocal<String> objectThreadLocal = newThreadLocal();
        objectThreadLocal.set("aaa");
        System.out.println(objectThreadLocal.get());
        MDC.put("traceId", "121211");
        Runnable runnable = new Runnable() {
            @Override
            public void run() {
                System.out.println(Thread.currentThread().getName() + ":" + objectThreadLocal.get() + ":" + MDC.get("traceId"));
            }
        };
//        ThreadTaskUtils.execute(runnable);
        GlobalThreadPool.submit(runnable);
    }

    public static <T> ThreadLocal<T> newThreadLocal() {
        return new TransmittableThreadLocal<T>();
    }

    public static <T> ThreadLocal<T> newThreadLocal(Supplier<T> supplier) {
        return new TransmittableThreadLocal<T>() {
            @Override
            protected T initialValue() {
                return supplier.get();
            }
        };
    }

    public static void sleepMilliSeconds(int time) {
        try {
            TimeUnit.MILLISECONDS.sleep(time);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 提供全局的默认任务执行线程池
     */
    private static ExecutorService GlobalThreadPool = newFixedThreadPool(Runtime.getRuntime().availableProcessors());


    /**
     * @param runnable
     */
    public static void submit(Runnable runnable) {
        GlobalThreadPool.submit(getRunnableWrap(runnable));
    }


    public static void execute(Runnable runnable) {
        GlobalThreadPool.execute(getRunnableWrap(runnable));
    }

    public static <T> Future<T> submit(Callable<T> callable) {
        Future<T> f = GlobalThreadPool.submit(getCallableWrap(callable));
        return f;
    }

    public static Runnable get(Runnable runnable) {
        return TtlRunnable.get(getRunnableWrap(runnable));
    }

    public static <T> Callable<T> get(Callable<T> callable) {
        return TtlCallable.get(getCallableWrap(callable));
    }

    private static <T> Callable<T> getCallableWrap(Callable<T> callable) {
        return isCallableWrap(callable) ? callable : new CallableWrap(callable);
    }

    private static Runnable getRunnableWrap(Runnable runnable) {
        return isRunnableWrap(runnable) ? runnable : new RunnableWrap(runnable);
    }

    /**
     * 返回固定线程数据的线程池，工作队列默认采用LinkedBlockingQueue，最大1024
     *
     * @param nThreads
     * @return
     */
    public static ExecutorService newFixedThreadPool(int nThreads) {
//        ThreadPoolExecutor threadPoolExecutor = ExecutorBuilder.create()
//                .setCorePoolSize(nThreads)
//                .setMaxPoolSize(nThreads)
//                .setKeepAliveTime(0L, TimeUnit.MILLISECONDS)
//                .setWorkQueue(new LinkedBlockingQueue(1024))
//                .build();
        ThreadPoolExecutor threadPoolExecutor = new MyThreadPoolExecutor(nThreads, nThreads, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(1024));
        ExecutorService ttlExecutorService = TtlExecutors.getTtlExecutorService(threadPoolExecutor);
        return ttlExecutorService;
    }

    private static class MyThreadPoolExecutor extends ThreadPoolExecutor {

        private ThreadPoolExecutor threadPoolExecutor;


        public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
        }

        public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
        }

        public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
        }

        public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        }

        @Override
        public Future<?> submit(Runnable task) {
            RunnableWrap newTask = new RunnableWrap(task);
            return super.submit(newTask);
        }

        @Override
        public <T> Future<T> submit(Callable<T> task) {
            CallableWrap<T> newTask = new CallableWrap(task);
            return super.submit(newTask);
        }

        @Override
        public <T> Future<T> submit(Runnable task, T result) {
            RunnableWrap newTask = new RunnableWrap(task);
            return super.submit(newTask, result);
        }

        @Override
        public void execute(Runnable command) {
            RunnableWrap newCommand = new RunnableWrap(command);
            super.execute(newCommand);
        }
    }

    /**
     * @param nThreads
     * @param workQueue 任务队列,使用时需要防止内存溢出
     * @return
     */
    public static ExecutorService newFixedThreadPool(int nThreads, BlockingQueue workQueue) {
//        ThreadPoolExecutor threadPoolExecutor = ExecutorBuilder.create()
//                .setCorePoolSize(nThreads)
//                .setMaxPoolSize(nThreads)
//                .setKeepAliveTime(0L, TimeUnit.MILLISECONDS)
//                .setWorkQueue(workQueue)
//                .setThreadFactory(Executors.defaultThreadFactory())
//                .setHandler(new ThreadPoolExecutor.AbortPolicy())
//                .build();
        ThreadPoolExecutor threadPoolExecutor = new MyThreadPoolExecutor(nThreads, nThreads, 0L, TimeUnit.MILLISECONDS, workQueue);
        ExecutorService ttlExecutorService = TtlExecutors.getTtlExecutorService(threadPoolExecutor);
        return ttlExecutorService;
    }

    /**
     * @param coreThreads   核心线程数
     * @param maxThreads    最大线程数
     * @param keepAliveTime 线程空闲时间
     * @return
     */
    public static ExecutorService newDynamicThreadPool(int coreThreads, int maxThreads, int keepAliveTime) {
//        ThreadPoolExecutor threadPoolExecutor = ExecutorBuilder.create()
//                .setCorePoolSize(coreThreads)
//                .setMaxPoolSize(maxThreads)
//                .setKeepAliveTime(keepAliveTime, TimeUnit.MILLISECONDS)
//                .setWorkQueue(new LinkedBlockingQueue(DEFAULT_WORK_QUEUE_CAPACITY))
//                .build();
        ThreadPoolExecutor threadPoolExecutor = new MyThreadPoolExecutor(coreThreads, maxThreads, keepAliveTime, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(DEFAULT_WORK_QUEUE_CAPACITY));
        return TtlExecutors.getTtlExecutorService(threadPoolExecutor);
    }

//    /**
//     * 提供原始的构造方式，谨慎使用
//     *
//     * @param builder
//     * @return
//     */
//    public static ExecutorService newFixedThreadPool(ExecutorBuilder builder) {
//
//        ThreadPoolExecutor threadPoolExecutor = builder.build();
//        ExecutorService ttlExecutorService = TtlExecutors.getTtlExecutorService(threadPoolExecutor);
//        return ttlExecutorService;
//    }

    static class RunnableWrap implements Runnable {
        private Map<String, String> mdcMap;
        private Runnable run;

        public RunnableWrap(Runnable runnable) {
            mdcMap = MDC.getCopyOfContextMap();
            this.run = runnable;
        }


        @Override
        public void run() {
            if (mdcMap != null) {
                Set<Map.Entry<String, String>> entries = mdcMap.entrySet();
                for (Map.Entry<String, String> e : entries) {
                    MDC.put(e.getKey(), e.getValue());
                }
                mdcMap = null;
            }
            run.run();
        }
    }

    public static boolean isRunnableWrap(Runnable o) {
        return o instanceof RunnableWrap;
    }

    public static boolean isCallableWrap(Callable o) {
        return o instanceof CallableWrap;
    }

    static class CallableWrap<T> implements Callable<T> {
        private Map<String, String> mdcMap;
        private Callable<T> callable;

        public CallableWrap(Callable<T> callable) {
            mdcMap = MDC.getCopyOfContextMap();
            this.callable = callable;
        }


        @Override
        public T call() throws Exception {
            if (mdcMap != null) {
                Set<Map.Entry<String, String>> entries = mdcMap.entrySet();
                for (Map.Entry<String, String> e : entries) {
                    MDC.put(e.getKey(), e.getValue());
                }
                mdcMap = null;
            }
            return callable.call();
        }
    }

}
